import torch

from sko.operators_gpu.random_bit_mask import random_bit_mask


def mutation(self):
    '''
    mutation of 0/1 type chromosome
    faster than `self.Chrom = (mask + self.Chrom) % 2`
    :param self:
    :return:
    '''
    mask = random_bit_mask(self.device, self.Lind, self.number_type, self.gpu_number_type, self.prob_mut, self.size_pop,
                           self.len_chrom, 3)
    self.Chrom ^= mask
    return self.Chrom


def mutation_vrptw(self):
    enable_mutation = torch.rand(self.size_pop, device=self.device) < self.prob_mut
    sel = torch.randint(0, 4, (self.size_pop,), device=self.device)

    self.tmpChrom[sel < 3 & enable_mutation] = swap(self.tmpChrom[sel < 3 & enable_mutation])

    self.tmpChrom_index[sel == 3 & enable_mutation] = change_maks(self.tmpChrom_index[sel == 3 & enable_mutation])

    self.Chrom[:] = 0
    self.Chrom[self.tmpChrom_index] = self.tmpChrom.reshape(-1)

    return self.Chrom


def change_maks(individual):
    a = torch.max((~individual[:, 1:]) * torch.randint(1, 65536, individual.shape, device=individual.device)[:, 1:],
                  dim=-1).indices + 1

    c = torch.randint(0, 2, (individual.shape[0],), device=individual.device)
    a = torch.where(c == 0, a, a - 1)
    a = torch.clamp(a, 0, individual.shape[-1] - 2)

    indx = torch.arange(individual.shape[0], device=individual.device)

    individual[indx, a], individual[indx, a + 1] = individual[indx, a + 1], individual[indx, a]

    return individual


def swap(individual):
    n1, n2 = torch.randint(0, individual.shape[-1], (2, individual.shape[0]), device=individual.device)
    n1, n2 = torch.minimum(n1, n2), torch.maximum(n1, n2)

    indx = torch.arange(individual.shape[0], device=individual.device)

    individual[indx, n1], individual[indx, n2] = individual[indx, n2], individual[indx, n1]

    return individual
